"""
@Date: 2022/3/16
@Author: pyn
"""
import gym
import numpy as np
from abc import abstractmethod
from gym.spaces import Space, Box, Discrete, Dict, MultiDiscrete
from wtpy.WtBtEngine import WtBtEngine, EngineType
from wtpy.StrategyDefs import BaseCtaStrategy, CtaContext, HftContext
from wtpy.ExtModuleDefs import BaseExtDataLoader
from typing import Union, Type, Tuple


class ActionPlan:
    """混入BaseCtaStrategy"""

    @abstractmethod
    def save_action(self, action):
        """保存环境动作到策略"""

    @property
    @abstractmethod
    def action_space(self) -> Space:
        """强化学习动作空间"""

    @property
    @abstractmethod
    def engine_type(self) -> EngineType:
        """"""
        raise NotImplementedError


class RewardPlan:
    """"""

    @abstractmethod
    def reset(self):
        """"""

    @abstractmethod
    def get_reward(self, context: CtaContext):
        """通过context计算reward"""

    @property
    @abstractmethod
    def reward(self) -> float:
        """返回当前时间步reward"""


class Observer:

    @property
    @abstractmethod
    def observation_space(self) -> Space:
        """强化学习观测空间"""

    @abstractmethod
    def get_obs(self, context: CtaContext):
        """通过context拿到内部观测"""

    @property
    @abstractmethod
    def obs(self) -> np.ndarray:
        """返回当前时间步的observation"""

    @abstractmethod
    def load_final_his_observations(self) -> bool:
        """从外部拿到每个时间步的obs并保存到实例内部"""

    @abstractmethod
    def pick_timerange(self) -> Tuple[int, int]:
        """环境reset时BtEngine挑选timerange策略"""


class Stopper:
    """"""

    @property
    @abstractmethod
    def is_stop(self) -> bool:
        """判断环境在该步是否结束"""


class Informer:
    """"""

    @property
    def info(self) -> dict:
        """输出环境在该步调试信息"""
        return {}


class WtpyCtaEnv(gym.Env):
    """"""

    def __init__(self,
                 extloader: BaseExtDataLoader,
                 Strategy: Union[Type[BaseCtaStrategy], Type[ActionPlan]],
                 rewardplan: RewardPlan,
                 observer: Observer,
                 stopper: Stopper,
                 informer: Informer = None,
                 slippage: int = 0):
        """"""
        self._iter_ = 0
        self._step_num_ = -1
        self._run_ = False

        self._engine = None

        self._Strategy = Strategy
        self._rewardplan = rewardplan
        self._observer = observer
        self._stopper = stopper
        self._informer = informer
        self._extloader = extloader

        self._slippage = slippage

    @property
    def action_space(self) -> Space:
        if hasattr(self, '_strategy'):
            return self._strategy.action_space

    @property
    def observation_space(self) -> Space:
        if hasattr(self, '_observer'):
            return self._observer.observation_space

    def reset(self):
        """"""
        self.close()
        self._iter_ += 1
        self._step_num_ += 1

        if self._engine is None:
            self._engine = WtBtEngine(eType=EngineType.ET_CTA, logCfg='./config/research/logcfgbt.yaml')
            self._engine.set_extended_data_loader(loader=self._extloader, bAutoTrans=True)
            # 根据etype选择
            self._engine.init('./config/common/', './config/research/configbt_cta.yaml')
            self._engine_step = self._engine.cta_step

            _timetange_ = self._observer.pick_timerange()
            s, e = _timetange_[0], _timetange_[-1]
            self._engine.configBacktest(stime=s, etime=e)
            self._engine.configBTStorage(mode='csv', path='./storage')
            self._engine.commitBTConfig()
        else:
            _timetange_ = self._observer.pick_timerange()
            s, e = _timetange_[0], _timetange_[-1]
            self._engine.set_time_range(beginTime=s, endTime=e)

        self._rewardplan.reset()

        self._strategy = self._Strategy(name='s1')
        self._strategy._env = self
        self._engine.set_cta_strategy(self._strategy, slippage=self._slippage, hook=True)

        self._engine.run_backtest(bAsync=True)
        self._run_ = True

        self._step()

        return self._observer.obs

    def step(self, action):
        """"""

        self._strategy.save_action(action)
        self._engine_step()  # on calc
        self._step()

        self._step_num_ += 1

        return self._observer.obs, self._rewardplan.reward, self._stopper.is_stop, self._informer.info

    def _step(self):
        """"""
        self._engine_step()  # on calc done
        # TODO 在这里检查是否触发timerange中endtime

    def close(self):
        """"""
        if self._engine is not None and self._run_:
            self._engine.stop_backtest()
            self._run_ = False

    def __del__(self):
        self._engine.release_backtest()
